"""
### 模块注释

#### 功能
该模块主要用于将PDF文件转换为纯文本文件，支持多线程处理和批量转换。核心功能包括：
1. **PDF页面图像提取**：使用`pdf2image`将PDF页面转换为图像。
2. **图像预处理**：提供`preprocess_image`函数用于图像预处理（当前为空实现，可根据需求扩展）。
3. **OCR文本识别**：使用`pytesseract`对图像进行OCR识别，提取文本。
4. **多线程处理**：支持多线程处理PDF页面，提高转换效率。
5. **批量转换**：支持递归遍历文件夹，批量转换所有PDF文件为文本文件。

#### 使用
1. **单文件转换**：
    ```python
    pdf_path = "path/to/your/pdf.pdf"
    output_text_path = "path/to/output.txt"
    pdf_to_text(pdf_path, output_text_path, use_multithreading=True)
    ```

2. **批量转换**：
    ```python
    source_dir = "path/to/source/directory"
    target_dir = "path/to/target/directory"
    batch_pdf_to_md(source_dir, target_dir, pdf_to_text)
    ```


#### 依赖安装
以下依赖需要提前安装：
1. **Python库**：
    ```bash
    pip install pdf2image pytesseract fitz Pillow
    ```

2. **Tesseract OCR**：
   - 下载并安装Tesseract OCR：[Tesseract官网](https://github.com/tesseract-ocr/tesseract)
   - 安装后，确保Tesseract可执行文件路径已添加到系统环境变量中。

#### 注意事项
1. **图像预处理**：当前`preprocess_image`函数为空，可根据需求添加灰度化、二值化、去噪等预处理操作。
2. **多线程配置**：`pdf_to_text`函数中的`thread_count`参数默认为CPU核心数的一半，可根据实际情况调整。
3. **OCR语言支持**：`lang`参数默认为`chi_sim+eng`（简体中文和英文），可根据需要修改为其他语言。

#### 示例
```python
if __name__ == "__main__":
    # 单文件转换测试
    _test_single_file()
    # 批量转换测试
    _test_batch_file()
```

"""
import os
import time

from PIL.Image import Image
from pdf2image import convert_from_path
import pytesseract
import fitz  # PyMuPDF
from concurrent.futures import ThreadPoolExecutor

from picture_to_text.utils import generate_file_pairs


def preprocess_image(image):
    pass
    # from PIL import Image, ImageEnhance, ImageFilter
    # 灰度化
    # image = image.convert('L')
    # # 二值化
    # image = image.point(lambda x: 0 if x < 128 else 255, '1')
    # # 去噪
    # image = image.filter(ImageFilter.MedianFilter())
    # # 确保图像为RGB模式
    # if image.mode != 'RGB':
    #     image = image.convert('RGB')
    # # 调整亮度和对比度
    # enhancer = ImageEnhance.Contrast(image)
    # image = enhancer.enhance(2.0)
    return image


#
#
# def process_page(page_num, pdf_path, lang):
#     images = convert_from_path(pdf_path, first_page=page_num + 1, last_page=page_num + 1, dpi=300, thread_count=8)
#     if images:
#         preprocessed_image = preprocess_image(images[0])
#         text = pytesseract.image_to_string(preprocessed_image, lang=lang, config='--psm 6 --oem 3')
#         text = '\n\n'.join([p.strip() for p in text.split('\n\n') if p.strip()])
#         return page_num, text
#     return page_num, ""
#
#
# def pdf_to_md_text_only(pdf_path, output_md, lang='chi_sim+eng', page_prefix="## ## #Page", use_multithreading=False,
#                         thread_pool_size=4, suffix="---"):
#     """
#     将PDF文件转换为仅包含文本的Markdown格式文件。
#
#     参数:
#         pdf_path (str): 输入PDF文件的路径。
#         output_md (str): 输出Markdown文件的路径。
#         lang (str, optional): 用于OCR识别的语言，默认为'chi_sim+eng'（简体中文和英文）。
#         page_prefix (str, optional): 每页内容的前缀，默认为"## ## #Page"。
#         use_multithreading (bool, optional): 是否启用多线程处理，默认为False。
#         thread_pool_size (int, optional): 线程池大小，仅在启用多线程时有效，默认为4。
#     """
#     doc = fitz.open(pdf_path)
#     md_content = [None] * len(doc)  # 预分配列表以保持顺序
#
#     if use_multithreading:
#         with ThreadPoolExecutor(max_workers=thread_pool_size) as executor:
#             futures = [executor.submit(process_page, page_num, pdf_path, lang) for page_num in range(len(doc))]
#
#             for future in futures:
#                 page_num, text = future.result()
#                 md_content[page_num] = f"{page_prefix} {page_num + 1}\n{text}\n\n{suffix}\n"
#     else:
#         for page_num in range(len(doc)):
#             text = ""
#             images = convert_from_path(pdf_path, first_page=page_num + 1, last_page=page_num + 1)
#             if images:
#                 preprocessed_image = preprocess_image(images[0])
#                 text = pytesseract.image_to_string(preprocessed_image, lang=lang, config='--psm 6 --oem 3')
#                 text = '\n\n'.join([p.strip() for p in text.split('\n\n') if p.strip()])
#             md_content[page_num] = f"{page_prefix} {page_num + 1}\n{text}\n\n{suffix}\n"
#
#     with open(output_md, "w", encoding="utf-8") as f:
#         f.write("\n".join(md_content))


def pdf_page_to_text(image: Image, lang='chi_sim+eng', config='--psm 6 --oem 3'):
    """
    将PDF页面图像转换为文本。

    该函数接收一个PDF页面的图像，使用Tesseract OCR引擎将其转换为文本。图像首先会经过预处理，
    然后使用指定的语言和配置进行文本识别。最后，对识别出的文本进行清理和格式化。

    参数:
    - image (Image): PDF页面的图像对象，通常为PIL.Image类型。
    - lang (str, 可选): Tesseract OCR引擎使用的语言模型，默认为'chi_sim+eng'，即简体中文和英文。
    - config (str, 可选): Tesseract OCR引擎的配置参数，默认为'--psm 6 --oem 3'，表示页面分割模式为6（假设为单一文本块），OCR引擎模式为3（默认LSTM引擎）。

    返回值:
    - str: 识别出的文本，经过清理和格式化后的字符串。
    """
    # 对图像进行预处理，以提高OCR识别的准确性
    preprocessed_image = preprocess_image(image)

    # 使用Tesseract OCR引擎将图像转换为文本
    text = pytesseract.image_to_string(preprocessed_image, lang=lang, config=config)

    # 清理和格式化文本，去除多余的空行和空白字符
    text = '\n\n'.join([p.strip() for p in text.split('\n\n') if p.strip()])

    return text


def _process_image(page_num, image: Image, start_page, lang, config):
    """
    多线程处理单张图像，提取文本并生成Markdown内容。
    """
    text = pdf_page_to_text(image, lang, config)
    return f"## ## #Page {start_page + page_num + 1}\n{text}\n\n---\n\n"


def pdf_to_text(pdf_path, output_text_path, batch_size=50, thread_count=None, use_multithreading=False,
                lang='chi_sim+eng', config='--psm 6 --oem 3'):
    """
    将PDF文件转换为文本文件，支持多线程处理。

    :param pdf_path: 输入PDF文件的路径。
    :param output_text_path: 输出文本文件的路径。
    :param batch_size: 每批处理的页数，默认为50。
    :param thread_count: 线程数，默认为CPU核心数的一半。
    :param use_multithreading: 是否启用多线程处理，默认为False。
    :param lang: 语言模型，默认为'chi_sim+eng
    :param config: 配置参数，默认为'--psm 6 --oem 3'，表示页面分割模式为6（假设为单一文本块），OCR引擎模式为3（默认LSTM引擎）。
    :return:
    """
    if thread_count is None:
        thread_count = int(os.cpu_count() / 2)
    if thread_count <= 0:
        raise ValueError("thread_count must be greater than 0")
    if batch_size <= 0:
        raise ValueError("batch_size must be greater than 0")
    with fitz.open(pdf_path) as doc:
        pages_length = len(doc)
        print(f"处理：{pdf_path} 总页数为：{pages_length}")
        for start_page in range(0, pages_length, batch_size):
            end_page = min(start_page + batch_size - 1, pages_length - 1)
            print(f"处理：{pdf_path} 第{start_page + 1}页到第{end_page + 1}页...")
            images = convert_from_path(
                pdf_path, first_page=start_page + 1, last_page=end_page + 1, dpi=300, thread_count=thread_count)
            md_content = ""
            if images:
                if use_multithreading:
                    with ThreadPoolExecutor(max_workers=thread_count) as executor:
                        futures = [executor.submit(
                            _process_image, page_num, image, start_page, lang=lang, config=config)
                            for page_num, image in enumerate(images)]
                        for future in futures:
                            md_content += future.result()
                else:
                    for page_num, image in enumerate(images):
                        md_content += _process_image(page_num, image, start_page, lang=lang, config=config)
            # 保存，首次覆盖，其余追加
            mode = 'w' if start_page == 0 else 'a'
            with open(output_text_path, mode, encoding='utf-8') as f:
                f.write(md_content)



def batch_pdf_to_md(source_dir, target_dir, pdf_to_text_func, search_str=".pdf", target_prefix=".md"):
    """
    递归遍历源文件夹，将名称包含指定字符串的所有PDF文件转换为Markdown文件，并保存到目标文件夹中，保持源文件的目录结构。

    :param target_prefix: 目标文件名前缀，默认为".md"。
    :param source_dir: 源文件夹路径。
    :param target_dir: 目标文件夹路径。
    :param pdf_to_text_func: 用于将PDF文件转换为Markdown文件的函数。
            该函数需要满足以下要求：
            - 参数1: pdf_path (str): 输入的PDF文件路径。
            - 参数2: output_md (str): 输出的Markdown文件路径。
            - 返回值: 无。该函数需要将转换结果直接写入到指定的Markdown文件中。
    :type pdf_to_text_func: function
    :param search_str: 要搜索的文件名包含的字符串，默认为".pdf"。
    :return: 无。该函数直接将结果写入到指定的目标文件夹中。
    """
    time_0 = time.time()
    i = 0
    for source_file_path, target_file_path in generate_file_pairs(source_dir, target_dir, search_str, target_prefix):
        # 调用传入的pdf_to_md_func函数进行转换
        current_time = time.time()
        print(f"< 第{i + 1}个 >\n开始转换：{source_file_path} -> {target_file_path}")
        pdf_to_text_func(source_file_path, target_file_path)
        print(f"< 第{i + 1}个 >\n转换完成：{source_file_path} -> {target_file_path}")
        print(f"转换耗时：{time.time() - current_time}秒\n总时间：{time.time() - time_0}秒\n\n")
        i += 1


def _test_single_file():
    # 使用示例
    pdf_path = r"C:\Users\Msi\Desktop\temp\PDF__txt_test\500_kV进线构架的抗震性能评估_商文念.pdf"
    # output_md = r"C:\Users\Msi\Desktop\temp\PDF__txt_test\sin_500_kV进线构架的抗震性能评估_商文念.md"
    # time_0 = time.time()
    # pdf_to_md_text_only(pdf_path, output_md)
    # print(f"单线程转换完成，耗时：{time.time() - time_0}秒")
    output_md = r"C:\Users\Msi\Desktop\temp\PDF__txt_test\mut_500_kV进线构架的抗震性能评估_商文念.md"
    time_0 = time.time()
    pdf_to_text(pdf_path, output_md, use_multithreading=True)
    print(f"多线程转换完成，耗时：{time.time() - time_0}秒")


def _test_batch_file():
    # 配置pdf_to_md_text_only的设置
    def custom_pdf_to_md(pdf_path, output_md):
        pdf_to_text(pdf_path, output_md, use_multithreading=True)

    source_dir = r"G:\works\books\abaqus\Abaqus2016中文帮助文档"
    target_dir = r"G:\works\books\abaqus_md\Abaqus2016中文帮助文档"
    batch_pdf_to_md(source_dir, target_dir, custom_pdf_to_md)


if __name__ == "__main__":
    print("pdf to txt 测试开始!!!")
    # _test_single_file()
    _test_batch_file()
